"""
利用SVM实现水果分类
数据集：爬虫从搜索结果中爬取的水果图像，
包含5个类别，共1036张（苹果288张、香蕉275张、葡萄216张、橙子276张、梨251张）
"""

import os
from typing import Union

import numpy as np
import cv2
import sklearn.metrics as sm
import sklearn.preprocessing as sp
from sklearn.model_selection import train_test_split
import sklearn.svm as svm
import pickle


class Data:
    def __init__(self, data_list):
        images, labels = data_list.T
        self.images = np.array([self.get_features(img) for img in images])
        self.labels = labels.astype('int64')

    @staticmethod
    def get_sift(im_gray):
        """
        使用sift算法获取图像的特征向量
        :param im_gray: 灰度图像的cv数据
        :return: 返回图片特征向量
        """
        h, w = im_gray.shape[:2]  # 取图像高度、宽度
        f = 200 / min(h, w)  # 计算缩放比率
        im_resize = cv2.resize(im_gray, None, fx=f, fy=f)  # 执行缩放

        # 计算图像特征
        sift = cv2.xfeatures2d.SIFT_create()
        keypoints = sift.detect(im_resize)  # 检测图像关键点
        _, desc = sift.compute(im_resize, keypoints)  # 计算特征

        desc = np.sum(desc, axis=0)  # 跨行求和

        return desc

    @staticmethod
    def get_features(img_path):
        """
        根据图片路径获取图片特征数据
        :param img_path: 图片路径
        :return: 返回图片特征向量
        """
        # 图像读取
        im_gray = cv2.imread(img_path, 0)  # 读取图像数据 0-灰度图像

        return Data.get_sift(im_gray)


class FruitData:
    def __init__(self, data_dir):
        self.name_dict = {'apple': 0, 'banana': 1, 'grape': 2, 'orange': 3, 'pear': 4}
        current_dir = os.path.dirname(__file__)
        self.data_dir = os.path.join(current_dir, data_dir)
        self.fruit_list = self.get_fruit_list(data_dir)
        train, test = train_test_split(np.array(self.fruit_list), test_size=0.1)
        self.train = Data(train)
        self.test = Data(test)

    def get_fruit_list(self, data_dir):
        fruit_dirs = os.listdir(self.data_dir)
        fruit_list = []
        for fruit_name in fruit_dirs:
            fruit_dir = os.path.join(self.data_dir, fruit_name)
            if not os.path.isdir(fruit_dir):
                continue
            imgs = os.listdir(fruit_dir)
            for img in imgs:
                img_path = os.path.join(fruit_dir, img)
                item = (img_path, self.name_dict[fruit_name])
                fruit_list.append(item)
        return fruit_list


class FruitModel:
    def __init__(self, data_dir, model_file=None):
        self.name_list = ['苹果', '香蕉', '葡萄', '橘子', '梨子']
        current_dir = os.path.dirname(__file__)
        self.data_dir = os.path.join(current_dir, data_dir)
        self.model_file = model_file and os.path.join(current_dir, model_file)
        self.model = self.get_model()
        self.data = None  # type:Union[FruitData,None]

    def get_model(self, load_model=True) -> svm.SVC:
        if self.model_file and load_model:
            if os.path.exists(self.model_file):
                with open(self.model_file, 'rb') as f:
                    model = pickle.load(f)
                return model
            else:
                return svm.SVC(kernel='poly', degree=2)  # 支持向量机模型
        else:
            return svm.SVC(kernel='poly', degree=2)  # 支持向量机模型

    def load_data(self):
        self.data = FruitData(self.data_dir)

    def train(self):
        print(self.model)
        self.model.fit(self.data.train.images, self.data.train.labels)
        print('训练结束')

    def save_model(self):
        if self.model_file and os.path.exists(os.path.dirname(self.model_file)):
            with open(self.model_file, 'wb') as f:
                pickle.dump(self.model, f)
                print("保存模型成功")
        else:
            print('保存模型失败')

    def eval(self):
        print('训练集分数:', self.model.score(self.data.train.images, self.data.train.labels))
        print('测试集分数:', self.model.score(self.data.test.images, self.data.test.labels))

        print('开始预测')
        pred_test = self.model.predict(self.data.test.images)
        print('预测结果', pred_test)
        print('r2预测精度:', sm.r2_score(self.data.test.labels, pred_test))
        print('=' * 30 + '分类报告' + '=' * 30)
        print(sm.classification_report(self.data.test.labels, pred_test))

    def play(self):
        while True:
            img_path = input('请输入图片路径：')
            if not img_path:
                break
            features = Data.get_features(img_path)
            pred = self.model.predict(features.reshape(1, -1))
            print('这是“%s”' % self.name_list[pred[0]])

    def recognize(self, img_bytes):
        """
        识别一张水果图片
        :param img_bytes: 图片的bytes数据
        :return: 识别的结果
        """
        img_array = np.fromstring(img_bytes, np.uint8)
        img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        features = Data.get_sift(img)
        pred = self.model.predict(features.reshape(1, -1))
        result = self.name_list[pred[0]]
        return result


fruit_model = FruitModel('fruits_data', 'AI_model/fruit.pkl')

if __name__ == '__main__':
    model = fruit_model

    print('begin loading data...')
    model.load_data()

    # print("begin training...")
    # model.train()
    # model.save_model()

    print('\nbegin testing...')
    model.eval()

    model.play()
